import sys
import os

from data.read_csv import read_csv_lines
from datetime import datetime
from data.image_embedding_milvus import add_infos

# 获取当前文件的目录
current_dir = os.path.dirname(os.path.abspath(__file__))
# 将当前文件的目录添加到 sys.path
sys.path.append(current_dir)


def embedding_image(img_infos, model_name, metric_type, user_id, batch_size=50):
    if user_id == '':
        print("请填写user_id")
        sys.exit(0)

    image_ids_list = img_infos["image_ids"]
    image_paths_list = img_infos["image_paths"]

    # 给加上list的每个元素加上相对路径
    image_paths_list = [item.replace("./", "./data/") for item in image_paths_list]

    print(image_ids_list)
    print(image_paths_list)

    # 记录开始时间
    start_time = datetime.now()
    print(f"开始时间: {start_time}")

    # 调用方法将图片进行向量化
    add_infos(image_ids_list, image_paths_list, user_id, batch_size, model_name, metric_type)

    # 记录结束时间
    end_time = datetime.now()
    print(f"结束时间: {end_time}")

    # 计算时间间隔
    duration = end_time - start_time
    # 打印时长
    print(f"执行时长: {duration}")


if __name__ == "__main__":
    file_path = 'data/csv/test_AI测试图片_1221-1.csv'
    model_name = "ViT-H-14"
    metric_type = "COSINE"
    user_id = '1000000'
    start_line = 1
    file_total_lines = 10000
    file_batch_step = 100
    emb_batch_size = 50

    for start in range(start_line, file_total_lines, file_batch_step):
        end = min(start + file_batch_step - 1, file_total_lines)
        img_infos = read_csv_lines(file_path, start, end)

        print("起始行数:", start, ". 结束行数：", end)
        embedding_image(img_infos, model_name, metric_type, user_id, emb_batch_size)

